from collections import Counter
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
from sklearn.datasets import load_iris

class KNN:
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None
        
    def fit(self, X, y):
        """存储训练数据"""
        self.X_train = X
        self.y_train = y
        return self
    
    def euclidean_distance(self, x1, x2):
        """计算欧几里得距离"""
        return np.sqrt(np.sum((x1 - x2) ** 2))
    
    def predict(self, X):
        """预测新数据的类别"""
        predictions = [self._predict(x) for x in X]
        return np.array(predictions)
    
    def _predict(self, x):
        """预测单个样本的类别"""
        # 计算所有训练样本的距离
        distances = [self.euclidean_distance(x, x_train) 
                    for x_train in self.X_train]
        
        # 获取最近的k个邻居的索引
        k_indices = np.argsort(distances)[:self.k]
        
        # 获取k个邻居的标签
        k_nearest_labels = [self.y_train[i] for i in k_indices]
        
        # 返回最常见的标签
        most_common = Counter(k_nearest_labels).most_common(1)
        return most_common[0][0]
    
    def predict_proba(self, X):
        """预测概率"""
        proba = []
        for x in X:
            distances = [self.euclidean_distance(x, x_train) 
                        for x_train in self.X_train]
            k_indices = np.argsort(distances)[:self.k]
            k_nearest_labels = [self.y_train[i] for i in k_indices]
            
            # 计算每个类别的概率
            counter = Counter(k_nearest_labels)
            proba_dict = {}
            total = len(k_nearest_labels)
            for label in set(self.y_train):
                proba_dict[label] = counter.get(label, 0) / total
            proba.append(proba_dict)
        
        return proba

# 测试KNN算法
def test_knn():
    # 加载鸢尾花数据集
    iris = load_iris()
    X, y = iris.data, iris.target
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42, stratify=y
    )
    
    # 应用KNN
    knn = KNN(k=3)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    
    # 评估模型
    accuracy = accuracy_score(y_test, y_pred)
    print(f"KNN准确率: {accuracy:.4f}")
    print("\n分类报告:")
    print(classification_report(y_test, y_pred, 
                              target_names=iris.target_names))
    
    # 可视化特征重要性（基于特征范围）
    feature_importance = np.std(X_train, axis=0)
    plt.figure(figsize=(10, 6))
    plt.bar(range(len(feature_importance)), feature_importance)
    plt.xticks(range(len(feature_importance)), iris.feature_names)
    plt.title('特征重要性（基于标准差）')
    plt.ylabel('标准差')
    plt.show()
    
    return knn, accuracy

# 运行KNN测试
knn_model, accuracy = test_knn()